/*
 * PyEigenVector.cpp
 *
 *  Created on: Jan 2, 2012
 *      Author: allan
 */

#include "PyEigenVector.h"

// Eigen includes
#include <Eigen/Core>
using namespace Eigen;

// Boost includes
#include <boost/python.hpp>
#include <boost/range.hpp>
#include <boost/shared_ptr.hpp>
#include <boost/python/slice.hpp>
using namespace boost::python;

// PyGeoReact includes
#define GEOREACT_NUMPY_INTERNAL
#include <Common/NumpyArray.h>
#include <Common/PyErrors.h>

template <typename Scalar, int NumpyType>
struct WrapperEigenVector
{
	typedef Matrix<Scalar, -1, 1> VectorType;
	
	static boost::shared_ptr<VectorType>
	init_default()
	{
		return boost::shared_ptr<VectorType>(new VectorType());
	}
	
	static boost::shared_ptr<VectorType>
	init_copy(const VectorType& copy)
	{
		return boost::shared_ptr<VectorType>(new VectorType(copy));
	}
	
	static boost::shared_ptr<VectorType>
	init_with_rows(int rows)
	{
		return boost::shared_ptr<VectorType>(new VectorType(rows));
	}
	
	static boost::shared_ptr<VectorType>
	init_with_rows_and_value(int rows, const Scalar& val)
	{
		boost::shared_ptr<VectorType> vec = init_with_rows(rows);
		
		std::fill_n(vec->data(), rows, val);
		
		return vec;
	}
	
	static boost::shared_ptr<VectorType>
	init_with_array(const numeric::array& array)
	{
		long rows = len(array);
		
		boost::shared_ptr<VectorType> vec = init_with_rows(rows);
		
		if(PyArray_TYPE(array.ptr()) != NumpyType)
			TypeError("type mismatch between array and vector entries");
		
		Scalar* data = (Scalar*)PyArray_DATA(array.ptr());
		
		std::copy(data, data + rows, vec->data());
		
		return vec;
	}
	
	template <typename Sequence>
	static boost::shared_ptr<VectorType>
	init_with_sequence(const Sequence& seq)
	{
		return init_with_array(numeric::array(seq));
	}
	
	static void
	get_slice_data(VectorType& self, const slice& slice, long& start, long& stop)
	{
		start = 0;
		stop  = self.rows();
		
		if(!slice.start().is_none())
			start = extract<long>(slice.start());
		
		if(!slice.stop().is_none())
			stop = extract<long>(slice.stop());
	
		if(start < 0 || stop < start || stop > self.rows())
			IndexError();
	}
	
	static void
	set_item(VectorType& self, int i, const Scalar& scalar)
	{
		self[i] = scalar;
	}
	
	static Scalar&
	get_item(VectorType& self, int i)
	{
		return self[i];
	}
	
	static void
	set_slice_with_scalar(VectorType& self, const slice& slice, const Scalar& scalar)
	{
		long start, stop;
		
		get_slice_data(self, slice, start, stop);
		
		for(int i = start; i < stop; ++i) 
			self[i] = scalar;
	}
	
	static void
	set_slice_with_eigen_vector(VectorType& self, const slice& slice, const VectorType& vec)
	{
		long start, stop;
		
		get_slice_data(self, slice, start, stop);
		
		const long rows = stop - start;
		
		if(vec.rows() != rows)
			IndexError();
		
		self.segment(start, rows) = vec;
	}
	
	static void
	set_slice_with_array(VectorType& self, const slice& slice, const numeric::array& array)
	{
		long start, stop;
		
		get_slice_data(self, slice, start, stop);
		
		long rows = stop - start;
		
		auto self_data = self.segment(start, rows).data();
		
		if(PyArray_TYPE(array.ptr()) != NumpyType)
			TypeError("type mismatch between array and vector entries");
		
		Scalar* data = (Scalar*)PyArray_DATA(array.ptr());
		
		std::copy(data, data + rows, self_data);
	}
	
	template <typename Sequence>
	static void
	set_slice_with_sequence(VectorType& self, const slice& slice, const Sequence& seq)
	{
		set_slice_with_array(self, slice, numeric::array(seq));
	}
	
	static object
	get_slice(VectorType& self, const slice& slice)
	{
		long start, stop;
		
		get_slice_data(self, slice, start, stop);
		
		const long rows = stop - start;
		
		return object(VectorType(self.segment(start, rows)));
	}
	
	static numeric::array 
	array(VectorType& self)
	{
		npy_intp dims[1] = { self.rows() };
		
		object obj(handle<>(
			incref(PyArray_SimpleNewFromData(1, dims, NumpyType, self.data()))));
		
		return extract<numeric::array>(obj);
	}
	
	static Scalar*
	begin(VectorType& self)
	{
		return self.data();
	}
	
	static Scalar*
	end(VectorType& self)
  	{
		return self.data() + self.rows() * self.cols();
	}
};

template <typename Scalar, int NumpyType>
void export_EigenVectorType(const char* class_name)
{
	typedef Matrix<Scalar, -1, 1> VectorType;
	typedef WrapperEigenVector<Scalar, NumpyType> WrapperType;
	
	class_<VectorType, boost::shared_ptr<VectorType>>(class_name)
		.def("__init__",
			make_constructor(&WrapperType::init_default))
		
		.def("__init__",
			make_constructor(&WrapperType::init_copy))
		
		.def("__init__",
			make_constructor(&WrapperType::init_with_rows))
		
		.def("__init__",
			make_constructor(&WrapperType::init_with_rows_and_value))
		
		.def("__init__",
			make_constructor(&WrapperType::init_with_array))
			
		.def("__init__",
			make_constructor(&WrapperType::template init_with_sequence<list>))
			
		.def("__init__",
			make_constructor(&WrapperType::template init_with_sequence<tuple>))
		
		.def("__len__",
			&VectorXd::rows)
		
		.def("__setitem__",
			&WrapperType::set_item,
				with_custodian_and_ward<1,2>()) // to let container keep value
		
		.def("__getitem__",
			&WrapperType::get_item,
				return_value_policy<copy_non_const_reference>())
		
		.def("__setitem__",
			&WrapperType::set_slice_with_scalar)
		
		.def("__setitem__",
			&WrapperType::set_slice_with_eigen_vector)
		
		.def("__setitem__",
			&WrapperType::set_slice_with_array)
			
		.def("__setitem__",
			&WrapperType::template set_slice_with_sequence<list>)
			
		.def("__setitem__",
			&WrapperType::template set_slice_with_sequence<tuple>)

		.def("__getitem__",
			&WrapperType::get_slice)
				
		.def("__iter__",
			range(&WrapperType::begin, &WrapperType::end))
			
		.def("__array__",
			&WrapperType::array)
			
		.add_property("rows",
			&VectorXd::rows)
		
		.def("array",
			&WrapperType::array)
		
		.def(self_ns::str(self_ns::self));
}

void export_EigenVector()
{
	export_EigenVectorType<double, NPY_DOUBLE>("VectorXd");
	export_EigenVectorType<float, NPY_FLOAT>("VectorXf");
	export_EigenVectorType<int, NPY_INT>("VectorXi");
}
